from data.dataset.factory import ComposeDataset, DatasetFactory
from easycore.common.registry import Registry

from .build import DatasetLoader


@DatasetLoader.register()
class PluginDataset:
    registry = Registry("plugin")

    @classmethod
    def register(cls, plugin_type_name=None, obj=None):
        return cls.registry.register(plugin_type_name, obj)

    @classmethod
    def get_plugins(cls, plugin_infos):
        plugins = [cls.registry(plugin_info['type'])(plugin_info) for plugin_info in plugin_infos]
        return plugins

    @classmethod
    def build(cls, dataset_config):
        dataset = DatasetLoader.get(dataset_config['dataset'])
        plugin_infos = dataset_config['plugins']

        plugins = cls.get_plugins(plugin_infos)
        for plugin in plugins:
            dataset = plugin(dataset)
        return dataset


@PluginDataset.register()
class RemoveNoObjectImage:
    def __init__(self, plugin_info):
        pass

    def __call__(self, dataset):
        return [data for data in dataset if data if len(data['objects']) > 0]


@PluginDataset.register()
class MapCategory:
    def __init__(self, plugin_info):
        self.category_map = plugin_info['category_map']

    def __call__(self, dataset):
        for data in dataset:
            new_objects = []
            for object in data['objects']:
                if object['category'] in self.category_map and self.category_map[object['category']] is not None:
                    new_objects.append(object)
            data['objects'] = new_objects
        return data


@PluginDataset.register()
class TrimBoxOutOfImage:
    def __init__(self, plugin_info):
        pass

    def __call__(self, dataset):
        for data in dataset:
            width = data['params']['width']
            height = data['params']['height']

            for object in data['objects']:
                box = object['box']
                object['box'] = [
                    min(0, box[0]),
                    max(0, box[1]),
                    min(width, box[2]),
                    max(height, box[3]),
                ]
        return dataset


@PluginDataset.register()
class RemoveSmallBoxObject:
    def __init__(self, plugin_info):
        # remove box s.t. the area of all box >= area_threshold
        self.area_threshold = plugin_info['area_threshold']

    def __call__(self, dataset):
        for data in dataset:
            new_objects = []
            for object in data['objects']:
                box = object['box']
                area = (box[2] - box[0]) * (box[3] - box[1])
                if area >= self.area_threshold:
                    new_objects.append(object)
            data['objects'] = new_objects
        return dataset


@PluginDataset.register()
class AddTransforms:
    def __init__(self, plugin_info):
        self.transform_configs = plugin_info['transforms']

    def __call__(self, dataset):
        for data in dataset:
            data['transforms'] += self.transform_configs
